import logging
import os
from abc import ABC, abstractmethod

import datasketch

import config
from tools import parseutility as parser
from tools.hashfunc import md5_hash32
from tools.utils import pearson_similarity
from rpc.ttypes import VulnResult, CVEInfo
from concurrent.futures import ThreadPoolExecutor

"""find的所有依赖"""


def query_cve_info(vuln_name: str, cve_query_cli) -> CVEInfo:
    """
    TODO：此函数是返回一个空内容的对象，后续若要支持根据漏洞名称（vuln_name）从漏洞查询客户端（cve_query_cli）
    查询漏洞相关数据，请重写本函数
    """
    return CVEInfo()


def generate_vuln_result(
        cve_query_cli,
        final_prob,
        func_feature: parser.FuncFeatures,
        vuln_func_feature: parser.FuncFeatures) -> VulnResult:
    return VulnResult(
        test_src_file=func_feature.get_feature('base_filename'),
        test_func=func_feature.get_feature('func_name'),
        prob=final_prob,
        from_repo_name=vuln_func_feature.get_feature('repo_name'),
        from_src_file=vuln_func_feature.get_feature('base_filename'),
        from_func=vuln_func_feature.get_feature('func_name'),
        vuln_fixed_commit=vuln_func_feature.get_feature('fixed_commit_id'),
        vuln_name=vuln_func_feature.get_feature('vuln_info'),
        cve_info=query_cve_info(vuln_func_feature.get_feature('vuln_info'), cve_query_cli)
    )


# TODO：当前给定的CVE_QUERY_CLI是空，后续若要实现漏洞的相关查询，请重写
CVE_QUERY_CLI = None


class VulnProbCalcStrategy(ABC):
    """
    抽象类：漏洞概率计算策略
    其子类：实现不同的具体计算策略
    """
    @abstractmethod
    def calc(self, init_prob, func_feature: parser.FuncFeatures, vuln_func_feature: parser.FuncFeatures):
        pass

    @staticmethod
    def create(*args, **kwargs):
        """静态工厂方法，提供默认概率计算策略"""
        return _DefaultVulnProbCalcStrategy()


class _DefaultVulnProbCalcStrategy(VulnProbCalcStrategy):
    """默认概率计算策略"""
    FEATURE_NAMES = [
        'func_LoC', 'param_num', 'lvar_num', 'dtype_num', 'funccall_num',
        'if_num', 'else_num', 'for_num', 'while_num', 'goto_num', 'switch_num', 'case_num', 'return_num'
    ]

    def __init__(self, encoding='utf-8'):
        self._changes_dict = {}
        self._encoding = encoding

    def calc(self, init_prob, func_feature: parser.FuncFeatures, vuln_func_feature: parser.FuncFeatures):
        """
        计算漏洞概率
        @param init_prob: 检测函数与漏洞函数之间原型相似度
        @param func_feature: 检测函数的函数特征
        @param vuln_func_feature: 漏洞函数的函数特征
        @return: 检测函数是漏洞函数的克隆概率值

        计算的基本策略为：
        0. 记init_prob为s0
        1. 计算检测函数与漏洞函数的相似度s1，记检测函数含有漏洞函数的修复属性比例为s2
        2. 最终概率final_prob=(a0*s0+a1*s1)*s2，其中ai为权重
        """
        vuln_func_abs_body = vuln_func_feature.get_feature('abs_func_body')
        if len(vuln_func_abs_body) <= 50:
            func_abs_body = func_feature.get_feature('abs_func_body')
            return 1.0 if vuln_func_abs_body == func_abs_body else 0.0

        parts_prob = {}

        feature_vec1 = []
        feature_vec2 = []
        for feature_name in self.FEATURE_NAMES:
            feature_vec1.append(func_feature.get_feature(feature_name))
            feature_vec2.append(vuln_func_feature.get_feature(feature_name))

        parts_prob['func_feature'] = pearson_similarity(feature_vec1, feature_vec2)  # 利用函数特征计算相似度
        parts_prob['func_feature'] = parts_prob['func_feature'] if parts_prob['func_feature'] >= 0 else 0
        parts_prob['abstract'] = init_prob  # 函数抽象对比相似度
        parts_prob['func_name'] = 1 if func_feature.get_feature('func_name').lower() == vuln_func_feature.get_feature(
            'func_name').lower() else 0

        if parts_prob['func_name']:
            weights = {
                'func_feature': 0.25,
                'abstract': 0.25,
                'func_name': 0.5,
            }
        else:
            weights = {
                'func_feature': 0.7,
                'abstract': 0.3,
                'func_name': 0,
            }

        final_prob = 0
        for key in parts_prob.keys():
            prob = parts_prob[key]
            final_prob += (weights[key] * prob)

        if final_prob > 1:
            final_prob = 1

        # 额外处理
        changes_file_fp = vuln_func_feature.get_feature('changes_info')
        if not changes_file_fp:
            raise ValueError('Not exists changes file.')
        changes_file_fp = os.path.join(config.ABS_BASE_DIR, changes_file_fp)

        diff_info = self._changes_dict.get(changes_file_fp, None)
        if not diff_info:
            deserializer = parser.PatchDiffInfo.get_deserializer()
            with open(changes_file_fp, 'r', encoding=self._encoding) as f:
                content = f.read().strip()
            diff_info: parser.PatchDiffInfo = deserializer.deserialize(content)
            self._changes_dict[changes_file_fp] = diff_info
        test_func_body = func_feature.get_feature('abs_func_body')
        add_shot = 0
        del_shot = 0

        def filter_(li):
            ret = []
            for _ in li:
                if len(_) > 10:
                    ret.append(_)
            return ret

        adds = filter_(diff_info.get_change_adds())  # fix漏洞新增的内容
        dels = filter_(diff_info.get_change_dels())     # fix漏洞删除的内容
        for add in adds:
            # fix漏洞新增的内容，原先没有
            if test_func_body.find(add) < 0:
                add_shot += 1

        for del_ in dels:
            if test_func_body.find(del_) >= 0:
                del_shot += 1

        adds_len = len(adds)
        dels_len = len(dels)

        prob_add = (final_prob * add_shot / adds_len) if adds_len > 0 else 0
        prob_del = (final_prob * del_shot / dels_len) if dels_len > 0 else 0
        return max(prob_add, prob_del)


class MinHashFactory(ABC):
    """
    抽象类：用于创建MinHash对象的工厂
    子类：实现本类
    """
    @abstractmethod
    def new_instance(self, func_feature: parser.FuncFeatures):
        pass

    @staticmethod
    def create(*args, **kwargs):
        """静态工厂方法，创建默认的MinHash工厂"""
        return _DefaultMinHashFactory(*args, **kwargs)


class _DefaultMinHashFactory(MinHashFactory):
    def __init__(self, hashfunc=md5_hash32, encoding=config.ENCODING):
        self._hashfunc = hashfunc
        self._encoding = encoding

    def new_instance(self, func_feature: parser.FuncFeatures):
        """
        创建min hash对象
        @param func_feature: 函数特征
        @return: min hash对象

        使用datasketch实现
        """
        m = datasketch.MinHash(hashfunc=self._hashfunc)
        abs_func_body = func_feature.get_feature('abs_func_body')
        b = [_.encode(self._encoding) for _ in parser.handled_func_split(abs_func_body)]
        m.update_batch(b)
        return m


class LSHFactory(ABC):
    @abstractmethod
    def new_instance(self):
        pass

    @staticmethod
    def create(*args, **kwargs):
        return _DefaultLSHFactory(*args, **kwargs)


class _DefaultLSHFactory(LSHFactory):
    def __init__(self, hashfunc=md5_hash32, threshold=0.3):
        self._hashfunc = hashfunc
        self._threshold = threshold

    def new_instance(self):
        return datasketch.MinHashLSH(threshold=self._threshold, hashfunc=self._hashfunc)


class FuncVulnContext(object):
    """
    封装了函数漏洞计算上下文，即对于给定的带检测函数特征，从漏洞函数数据集中检索出所有
    超过漏洞阈值的函数漏洞
    """

    def __init__(
            self,
            min_hash_factory: MinHashFactory = MinHashFactory.create(),
            lsh_factory: LSHFactory = LSHFactory.create(),
            vuln_prob_calc_strategy: VulnProbCalcStrategy = VulnProbCalcStrategy.create(),
            final_prob_threshold=0.7,
            **kwargs
    ):
        """
        init方法
        @param min_hash_factory: min hash工厂
        @param lsh_factory: lsh工厂
        @param vuln_prob_calc_strategy: 漏洞函数计算策略
        @param final_prob_threshold: 最终函数漏洞过滤阈值
        @param kwargs: 附加参数
        """
        lsh = lsh_factory.new_instance()
        if isinstance(lsh, datasketch.MinHashLSH):
            self._insert_func = self.__insert_min_hash_lsh
            self._query_func = self.__query_min_hash_lsh
        elif isinstance(lsh, datasketch.MinHashLSHForest):
            self._insert_func = self.__insert_min_hash_lsh_forest
            self._query_func = self.__query_min_hash_lsh_forest
            self._top_k = kwargs.get('top_k', 5)
        self._lsh = lsh
        self._dict = {}
        self._min_hash_factory = min_hash_factory
        self._vuln_prob_calc_strategy = vuln_prob_calc_strategy
        self._final_prob_threshold = final_prob_threshold

    def __insert_min_hash_lsh(self, key: str, min_hash: datasketch.MinHash):
        lsh: datasketch.MinHashLSH = self._lsh
        lsh.insert(key, min_hash)

    def __insert_min_hash_lsh_forest(self, key: str, min_hash: datasketch.MinHash):
        lsh: datasketch.MinHashLSHForest = self._lsh
        lsh.add(key, min_hash)

    def __query_min_hash_lsh(self, min_hash: datasketch.MinHash):
        lsh: datasketch.MinHashLSH = self._lsh
        return lsh.query(min_hash)

    def __query_min_hash_lsh_forest(self, min_hash: datasketch.MinHash):
        lsh: datasketch.MinHashLSHForest = self._lsh
        return lsh.query(min_hash, self._top_k)

    def insert_min_hash(self, key: str, func_feature: parser.FuncFeatures):
        min_hash_factory: MinHashFactory = self._min_hash_factory
        min_hash = min_hash_factory.new_instance(func_feature)
        self._insert_func(key, min_hash)
        self._dict[key] = (func_feature, min_hash)

    def _query_func_feature(self, key: str) -> parser.FuncFeatures:
        return self._dict[key][0]

    def _query_min_hash(self, key: str) -> datasketch.MinHash:
        return self._dict[key][1]

    def query(self, func_feature: parser.FuncFeatures) -> list:
        """
        即对于给定的带检测函数特征，从漏洞函数数据集中检索出所有超过漏洞阈值的函数漏洞
        @param func_feature: 带检测函数的函数特征
        @return:
        """
        min_hash_factory: MinHashFactory = self._min_hash_factory
        min_hash = min_hash_factory.new_instance(func_feature)
        rst = self._query_func(min_hash)
        ret = []
        if not rst:
            return ret
        vuln_prob_calc_strategy: VulnProbCalcStrategy = self._vuln_prob_calc_strategy
        for candidate_key in rst:
            init_prob = self.jaccard(min_hash, candidate_key)
            vuln_func_feature = self._query_func_feature(candidate_key)
            final_prob = vuln_prob_calc_strategy.calc(init_prob, func_feature, vuln_func_feature)
            if final_prob >= self._final_prob_threshold:
                ret.append(generate_vuln_result(CVE_QUERY_CLI, final_prob, func_feature, vuln_func_feature))
        return ret

    def jaccard(self, a, b):
        """查询两个集合的jaccard值"""
        def _adapt(x):
            if isinstance(x, str):
                return self._query_min_hash(x)
            elif isinstance(a, datasketch.MinHash):
                return x
            else:
                raise TypeError('Just supports type str and MinHash, but type %s.' % type(x))

        return _adapt(a).jaccard(_adapt(b))


class FuncVulnQueryContext(object):
    """
    封装了漏洞检测的上下文
    """
    def __init__(
            self,
            func_vuln_context: FuncVulnContext,
            pool: ThreadPoolExecutor,
            alv: int = 5,
            normalize: bool = True
    ):
        """
        init方法
        @param func_vuln_context: 函数漏洞计算上下文
        @param pool: 线程池
        @param alv: 函数原型抽象等级
        @param normalize: 函数原型是否标准化
        """
        self._load_vuln_func_features = False
        self._func_vuln_context = func_vuln_context
        self._pool = pool
        self._alv = alv,
        self._normalize = normalize

    def load_vuln_func_features(self, vuln_func_features_files: list, encoding='utf-8') -> int:
        """
        以给定的encoding文件编码格式从给定的vuln_func_features_files漏洞函数特征文件列表中
        加载所有的漏洞函数特征到内存

        @param vuln_func_features_files: 漏洞函数特征文件列表
        @param encoding: 文件编码格式
        @return: 返回加载的函数特征个数
        """
        deserializer: parser.Deserializer = parser.FuncFeatures.get_deserializer()

        def generate_key(func_feature: parser.FuncFeatures):
            return '_'.join(
                [
                    func_feature.get_feature('vuln_info'),
                    func_feature.get_feature('repo_name'),
                    func_feature.get_feature('fixed_commit_id'),
                    func_feature.get_feature('base_filename'),
                    func_feature.get_feature('func_name'),
                    func_feature.get_feature('abs_func_body_hash')
                ]
            )

        def handle_vuln_func(file):
            with open(file, 'r', encoding=encoding) as f:
                content = f.read().strip()
            feature_: parser.FuncFeatures = deserializer.deserialize(content)
            self._func_vuln_context.insert_min_hash(generate_key(feature_), feature_)

        for _ in self._pool.map(handle_vuln_func, vuln_func_features_files):
            pass
        self._load_vuln_func_features = True
        return len(vuln_func_features_files)

    def __check_load_vuln_func_features(self):
        if not self._load_vuln_func_features:
            raise RuntimeError('plz invoke the func load_vuln_func_features before calling query.')

    def query_file(self, src_file: str, encoding=config.ENCODING) -> tuple:
        return self.query_files(src_files=[src_file], encoding=encoding)

    def query_files(self, src_files: list, encoding=config.ENCODING,
                    do_after_per_file=None, do_after_per_file_func_args=None) -> tuple:
        """
        对于给定的所有待检测的源代码文件src_files，对其中所有的函数进行漏洞检测并返回检测结果
        @param src_files: 待检测的源代码文件列表
        @param encoding: 文件的编码格式
        @param do_after_per_file: 在检测完每个源文件后的后续动作函数
        @param do_after_per_file_func_args: 检测完每个源文件后的后续动作函数的函数参数
        @return: 检测结果
        """
        if do_after_per_file_func_args is None:
            do_after_per_file_func_args = []
        self.__check_load_vuln_func_features()

        def filter_test_func(test_func_feature: parser.FuncFeatures) -> bool:
            return len(test_func_feature.get_feature('abs_func_body')) <= 50

        def handle_test_func(test_file):
            ret = []
            lines = 0
            funcs = parser.DEFAULT_FUNC_PARSER.parse(filepath=test_file, shallow=False)
            logging.info('Finds %s needed-test funcs in src file %s', len(funcs), test_file)
            for feature_ in [parser.extract_features(_, 5) for _ in funcs]:
                file_name = feature_.get_feature('base_filename')
                func_name = feature_.get_feature('func_name')
                func_LoC = feature_.get_feature('func_LoC')
                lines += func_LoC
                if filter_test_func(feature_):
                    logging.info('The func named %s from %s triggered filtering condition, filtered.', func_name,
                                 file_name)
                    continue
                ret.extend(self._func_vuln_context.query(feature_))
            if do_after_per_file:
                if do_after_per_file_func_args:
                    do_after_per_file(*do_after_per_file_func_args)
                else:
                    do_after_per_file()
            return len(funcs), lines, ret

        ret_ = []
        test_func_num_ = 0
        test_lines_ = 0
        for funcs_num, lines, rst in self._pool.map(handle_test_func, src_files):
            test_func_num_ += funcs_num
            test_lines_ += lines
            ret_.extend(rst)
        return test_func_num_, test_lines_, ret_
